'''
sr:     flask_gfpgan_query.py       port 20018
pose:   flask_yolo_pose_query.py    port 20020
parse:  flask_schp_query.py         port 20017
densepose: flask_densepose_query.py port 20016

# 输入 img
# 分别获得 超分sr，
# pose data+绘制图，parsing data+图，densepose 图

img_dir = /mnt/nas/shengjie/datasets/reid_sysu30k/sysu_train_set_all_part1/0000000001
save_dir = /mnt/nas/shengjie/datasets/reid_sysu30k_lsj/sysu_train_set_all_part1/0000000001
    save name.jpg

获得sr图之后，再送入下面的进行处理

save_dir = /mnt/nas/shengjie/datasets/reid_sysu30k_lsj/sysu_train_set_all_part1_pose/0000000001
    save name.np and name.jpg
save_dir = /mnt/nas/shengjie/datasets/reid_sysu30k_lsj/sysu_train_set_all_part1_hp/0000000001
    save name.jpg
save_dir = /mnt/nas/shengjie/datasets/reid_sysu30k_lsj/sysu_train_set_all_part1_dp/0000000001
    save name_bg.jpg and name_nobg.jpg

'''
import os
import requests
import json
import base64
import numpy as np

def process_sr(img_path, save_path):
    """调用超分接口，处理图片并保存。
    成功返回(True, 保存的path)，失败返回(False, "")
    """
    url = 'http://127.0.0.1:20018/sr'
    with open(img_path, 'rb') as f:
        files = {'image': f}
        response = requests.post(url, files=files)
        if response.status_code == 200:
            with open(save_path, 'wb') as out:
                out.write(response.content)
            print(f'SR结果已保存: {save_path}')
            return True, save_path
        else:
            print(f'SR失败, 状态: {response.status_code}, {response.text}')
            return False, ""

def process_pose(img_path, save_dir, save_name):
    """调用姿态接口，获取关键点并画图，保存为npy和jpg
    成功返回(True, (npy_path, img_save_path)), 失败返回(False, "")
    """
    url = 'http://127.0.0.1:20020/api/get_pose_by_yolo_batch'
    response = requests.post(url, files=[('images', (os.path.basename(img_path), open(img_path, 'rb'), 'image/jpeg'))])
    if response.status_code == 200:
        result = response.json()
        kps = None
        npy_path = ""
        img_save_path = ""
        try:
            # 取0号人的关键点
            kps = result['keypoints'][0]
        except Exception as e:
            print('获取pose关键点失败:', e)
            return False, ""
        
        # 保存关键点
        os.makedirs(save_dir, exist_ok=True)
        npy_path = os.path.join(save_dir, save_name + '.npy')
        np.save(npy_path, np.array(kps))

        # 再绘图
        url_draw = 'http://127.0.0.1:20020/api/draw_pose'
        files2 = {'image': open(img_path, 'rb')}
        data = {
            'kps': str(kps),
            'threshold': 0.3,
            'radius': 5
        }
        resp2 = requests.post(url_draw, files=files2, data=data)
        if resp2.status_code == 200:
            img_save_path = os.path.join(save_dir, save_name + '.jpg')
            with open(img_save_path, 'wb') as f:
                f.write(resp2.content)
            print(f"pose绘制结果已保存: {img_save_path}")
            return True, (npy_path, img_save_path)
        else:
            print('pose绘制失败:', resp2.status_code)
            return False, ""
    else:
        print('pose接口失败:', response.status_code)
        return False, ""

def process_hp(img_path, save_dir, save_name):
    """调用SCHP parsing接口，保存npy和jpg
    成功返回(True, img_save_path)，失败返回(False, "")
    """
    url = 'http://127.0.0.1:20017/parse'
    with open(img_path, 'rb') as f:
        files = {'image': f}
        response = requests.post(url, files=files)
    if response.status_code == 200:
        os.makedirs(save_dir, exist_ok=True)
        img_save_path = os.path.join(save_dir, save_name + '.jpg')
        with open(img_save_path, 'wb') as out:
            out.write(response.content)
        print(f'parsing分割图已保存: {img_save_path}')
        # 假设接口暂不直接返回分割np数据，仅保存图片
        # 若后续接口加npy返回，再补充保存npy
        return True, img_save_path
    else:
        print('hp parsing接口请求失败:', response.status_code)
        return False, ""

def process_dp(img_path, save_dir, save_name):
    """调用densepose接口，保存bg/nobg图
    成功返回(True, (bg_path, nobg_path)), 失败返回(False, "")
    """
    url = 'http://127.0.0.1:20016/densepose'
    with open(img_path, 'rb') as f:
        files = {'image': f}
        response = requests.post(url, files=files)
    if response.status_code == 200:
        json_data = response.json()
        bg_base64_str = json_data.get("bg_result_png_base64", "")
        nobg_base64_str = json_data.get('nobg_result_png_base64',"")
        os.makedirs(save_dir, exist_ok=True)
        bg_path = ""
        nobg_path = ""
        if bg_base64_str:
            bg_path = os.path.join(save_dir, save_name+'_bg.jpg')
            img_bytes = base64.b64decode(bg_base64_str)
            with open(bg_path, 'wb') as out_img:
                out_img.write(img_bytes)
            print(f'DensePose带背景已保存: {bg_path}')
        if nobg_base64_str:
            nobg_path = os.path.join(save_dir, save_name+'_nobg.jpg')
            nobg_image_bytes = base64.b64decode(nobg_base64_str)
            with open(nobg_path,'wb') as f:
                f.write(nobg_image_bytes)
            print(f'DensePose无背景已保存: {nobg_path}')
        if bg_path or nobg_path:
            return True, (bg_path, nobg_path)
        else:
            return False, ""
    else:
        print(f'DensePose接口请求失败: {response.status_code}, {response.text}')
        return False, ""

def main_process(img_path):
    import os

    def safe_remove(paths):
        """Utility function to remove files if path is not empty."""
        for p in paths:
            if isinstance(p, (list, tuple)):
                for item in p:
                    if item and os.path.exists(item):
                        try:
                            os.remove(item)
                        except Exception as e:
                            print(f"删除失败 {item}: {e}")
            elif p and os.path.exists(p):
                try:
                    os.remove(p)
                except Exception as e:
                    print(f"删除失败 {p}: {e}")

    img_dir = os.path.dirname(img_path)
    base_name = os.path.splitext(os.path.basename(img_path))[0]
    saved_paths = []  # 用于保存本次处理所有新生成的文件路径

    # 1. 超分
    save_dir_sr = img_dir.replace("reid_sysu30k", "reid_sysu30k_lsj")
    os.makedirs(save_dir_sr, exist_ok=True)
    sr_path = os.path.join(save_dir_sr, base_name + '.jpg')
    sr_success = process_sr(img_path, sr_path)
    if isinstance(sr_success, tuple):
        sr_flag, sr_ret_path = sr_success
    else:
        sr_flag, sr_ret_path = sr_success, sr_path if sr_success else ""
    if not sr_flag or not sr_ret_path or not os.path.exists(sr_ret_path):
        print('SR step failed, abort.')
        safe_remove([sr_ret_path])
        return False
    saved_paths.append(sr_ret_path)

    # 2. pose
    save_dir_pose = save_dir_sr.replace("_lsj", "_lsj_pose")
    pose_flag, pose_ret_path = process_pose(sr_ret_path, save_dir_pose, base_name)
    if not pose_flag:
        print('Pose step failed, abort.')
        safe_remove(saved_paths + [pose_ret_path])
        return False
    saved_paths.append(pose_ret_path)

    # 3. hp parsing
    save_dir_hp = save_dir_sr.replace("_lsj", "_lsj_hp")
    hp_flag, hp_ret_path = process_hp(sr_ret_path, save_dir_hp, base_name)
    if not hp_flag or not hp_ret_path or not os.path.exists(hp_ret_path):
        print('HP Parsing step failed, abort.')
        safe_remove(saved_paths + [hp_ret_path])
        return False
    saved_paths.append(hp_ret_path)

    # 4. densepose
    save_dir_dp = save_dir_sr.replace("_lsj", "_lsj_dp")
    dp_flag, dp_ret = process_dp(sr_ret_path, save_dir_dp, base_name)
    # dp_ret may be tuple like (bg_path, nobg_path)
    paths_dp = []
    if isinstance(dp_ret, (list, tuple)):
        # 过滤掉空的和None的路径
        paths_dp = [p for p in dp_ret if p]
    elif isinstance(dp_ret, str):
        if dp_ret: paths_dp = [dp_ret]
    if not dp_flag or not paths_dp or not all([os.path.exists(p) for p in paths_dp]):
        print('DensePose step failed, abort.')
        safe_remove(saved_paths + paths_dp)
        return False
    saved_paths.extend(paths_dp)

    # 所有步骤成功
    return True

if __name__=='__main__':
    # img_dir = /mnt/nas/shengjie/datasets/reid_sysu30k/sysu_train_set_all_part1/0000000001
    # 遍历一个目录下的所有图片，批量处理
    import glob

    # 指定你要处理的文件夹
    img_dir = '/mnt/nas/shengjie/datasets/reid_sysu30k/sysu_train_set_all_part1/0000000001'
    img_paths = sorted(glob.glob(os.path.join(img_dir, '*.jpg')))
    print(f'发现{len(img_paths)}张图片，开始处理')
    failed_imgs = []
    from tqdm import tqdm
    for img_path in tqdm(img_paths):
        print(f'处理: {img_path}')
        success = main_process(img_path)
        if not success:
            failed_imgs.append(img_path)
    if failed_imgs:
        print(f"处理失败的文件如下：")
        for f in failed_imgs:
            print(f)
        # 将失败的文件保存到 img_dir/failed_process.txt
        failed_txt_path = os.path.join(img_dir, "failed_process.txt")
        with open(failed_txt_path, 'w') as f:
            for failed_img in failed_imgs:
                f.write(failed_img + "\n")
        print(f"已将失败文件列表保存到 {failed_txt_path}")
    else:
        print("全部文件处理成功。")